import torch
import torch.nn as nn

"""RNN 模块。"""


class RNNLayer(nn.Module):
    """
    循环神经网络层
    """
    def __init__(self, inputs_dim, outputs_dim, recurrent_N, use_orthogonal):
        """
        初始化RNN层
        
        参数:
            inputs_dim: 输入维度
            outputs_dim: 输出维度
            recurrent_N: 循环层数
            use_orthogonal: 是否使用正交初始化
        """
        super(RNNLayer, self).__init__()
        self._recurrent_N = recurrent_N  # 循环层数
        self._use_orthogonal = use_orthogonal  # 是否使用正交初始化

        # 创建GRU网络
        self.rnn = nn.GRU(inputs_dim, outputs_dim, num_layers=self._recurrent_N)
        # 初始化GRU参数
        for name, param in self.rnn.named_parameters():
            if 'bias' in name:
                nn.init.constant_(param, 0)  # 偏置初始化为0
            elif 'weight' in name:
                if self._use_orthogonal:
                    nn.init.orthogonal_(param)  # 使用正交初始化
                else:
                    nn.init.xavier_uniform_(param)  # 使用Xavier均匀初始化
        # 创建层归一化
        self.norm = nn.LayerNorm(outputs_dim)

    def forward(self, x, hxs, masks):
        """
        前向传播
        
        参数:
            x: 输入特征
            hxs: 隐藏状态
            masks: 掩码，表示哪些序列是有效的
            
        返回:
            x: 处理后的特征
            hxs: 更新后的隐藏状态
        """
        if x.size(0) == hxs.size(0):
            # 如果批次大小相同，直接处理
            x, hxs = self.rnn(x.unsqueeze(0),
                              (hxs * masks.repeat(1, self._recurrent_N).unsqueeze(-1)).transpose(0, 1).contiguous())
            x = x.squeeze(0)
            hxs = hxs.transpose(0, 1)
        else:
            # x是一个(T, N, -1)张量，已被展平为(T * N, -1)
            N = hxs.size(0)  # 批次大小
            T = int(x.size(0) / N)  # 序列长度

            # 重新整形
            x = x.view(T, N, x.size(1))

            # 对掩码做同样的处理
            masks = masks.view(T, N)

            # 找出序列中任何智能体有零的步骤
            # 我们总是假设t=0有一个零，因为这使逻辑更清晰
            has_zeros = ((masks[1:] == 0.0)
                         .any(dim=-1)
                         .nonzero()
                         .squeeze()
                         .cpu())

            # +1来修正masks[1:]
            if has_zeros.dim() == 0:
                # 处理标量
                has_zeros = [has_zeros.item() + 1]
            else:
                has_zeros = (has_zeros + 1).numpy().tolist()

            # 将t=0和t=T添加到列表中
            has_zeros = [0] + has_zeros + [T]

            hxs = hxs.transpose(0, 1)

            outputs = []
            for i in range(len(has_zeros) - 1):
                # 我们现在可以一起处理掩码中没有零的步骤！
                # 这样更快
                start_idx = has_zeros[i]  # 开始索引
                end_idx = has_zeros[i + 1]  # 结束索引
                # 准备隐藏状态
                temp = (hxs * masks[start_idx].view(1, -1, 1).repeat(self._recurrent_N, 1, 1)).contiguous()
                # 处理当前片段
                rnn_scores, hxs = self.rnn(x[start_idx:end_idx], temp)
                outputs.append(rnn_scores)

            # 断言输出长度等于T
            # x是一个(T, N, -1)张量
            x = torch.cat(outputs, dim=0)

            # 展平
            x = x.reshape(T * N, -1)
            hxs = hxs.transpose(0, 1)

        # 应用层归一化
        x = self.norm(x)
        return x, hxs
